import math
import oneflow as torch
import oneflow.nn as nn
import oneflow.nn.functional as F


class SlidingWindow(nn.Module):
    def __init__(self, model_size, left_context=5, right_context=0, stride=1):
        super(SlidingWindow, self).__init__()

        self.model_size = model_size
        self.left_contxt = left_context
        self.right_context = right_context
        self.chunk_size = left_context + right_context + 1
        # padding_len = math.floor(chunk_size / 2)
        self.unfold1 = torch.nn.Unfold(kernel_size=(self.chunk_size, self.model_size), padding=0, stride=stride)
        self.unfold2 = torch.nn.Unfold(kernel_size=(self.chunk_size, 1), padding=0, stride=stride)

    def forward(self, inputs, mask=None):
        # inputs: [n, t, v]
        # mask: [n, 1, t]
        b = inputs.size(0)

        inputs = F.pad(inputs, pad=(0, 0, self.left_contxt, self.right_context), value=0.0)
        inputs = inputs.unsqueeze(1)
        output = self.unfold1(inputs)
        nb = output.size(2)
        output = output.transpose(1, 2)
        chunk_output = output.reshape(b, nb, self.chunk_size, self.model_size)

        if mask is not None:
            mask = F.pad(mask.float(), pad=(self.left_contxt, self.right_context), value=0)
            mask = mask.unsqueeze(-1)       
            frame_mask = self.unfold2(mask)
            frame_mask = frame_mask.transpose(1, 2).gt(0)
            chunk_mask = torch.sum(frame_mask, dim=-1).gt(0)
        else:
            frame_mask = None
            chunk_mask = None

        return chunk_output, frame_mask, chunk_mask


# if __name__ == '__main__':

#     # cutter1 = SlidingWindow(4, 1, 2)
#     cutter2 = SlidingWindow(4, 3, 1, 2)
#     inputs = torch.rand([2, 5, 4])
#     mask = torch.LongTensor([[1, 1, 1, 1, 1], [1, 0, 0, 0, 0]]).unsqueeze(1) > 0
#     print(inputs)
#     print(mask)
#     output, frame_mask, chunk_mask = cutter2(inputs, mask)
#     # print(output)
#     # print(output.size())
#     print(frame_mask)
#     print(frame_mask.size())
#     print(chunk_mask)
#     print(chunk_mask.size())